import React, { useEffect, useState } from 'react';
import { useParams, useNavigate } from 'react-router-dom';
import { 
  Layout, 
  Card, 
  Typography, 
  Button, 
  Space, 
  Form, 
  Table, 
  Input, 
  Modal, 
  Divider, 
  message, 
  Badge,
  Descriptions,
  Progress
} from 'antd';
import { ArrowLeftOutlined, PlusOutlined, DeleteOutlined, PlayCircleOutlined, PauseCircleOutlined, SaveOutlined } from '@ant-design/icons';
import { Agent, getAgentById, updateAgentSystemPrompt } from '@/api/agent';
import { 
  Training, 
  TrainingExample, 
  TrainingSession,
  TrainingLog,
  createTraining, 
  startTraining, 
  applyTrainingResult,
  getTrainingsByAgentId,
  createExample,
  getExamplesByTrainingId,
  deleteExample,
  getSessionsByTrainingId,
  getLogsBySessionId
} from '@/api/training';
import './style.less';

const { Title, Text, Paragraph } = Typography;
const { Content, Sider } = Layout;
const { TextArea } = Input;

interface TrainingFormValues {
  name: string;
  description: string;
  maxIterations: number;
}

const AgentTrain: React.FC = () => {
  const { id } = useParams<{ id: string }>();
  const navigate = useNavigate();
  const [agent, setAgent] = useState<Agent | null>(null);
  const [loading, setLoading] = useState<boolean>(true);
  const [trainings, setTrainings] = useState<Training[]>([]);
  const [currentTraining, setCurrentTraining] = useState<Training | null>(null);
  const [examples, setExamples] = useState<TrainingExample[]>([]);
  const [isTrainingModalVisible, setIsTrainingModalVisible] = useState<boolean>(false);
  const [isExampleModalVisible, setIsExampleModalVisible] = useState<boolean>(false);
  const [form] = Form.useForm<TrainingFormValues>();
  const [exampleForm] = Form.useForm<TrainingExample>();
  const [trainingSessions, setTrainingSessions] = useState<TrainingSession[]>([]);
  const [trainingLogs, setTrainingLogs] = useState<TrainingLog[]>([]);
  const [isLoading, setIsLoading] = useState<boolean>(false);

  useEffect(() => {
    if (id) {
      fetchAgentData();
      fetchTrainings();
    }
  }, [id]);

  useEffect(() => {
    if (currentTraining) {
      fetchExamples();
      fetchSessions();
    }
  }, [currentTraining]);

  const fetchAgentData = async () => {
    try {
      setLoading(true);
      const data = await getAgentById(Number(id));
      setAgent(data);
    } catch (error) {
      message.error('获取智能体信息失败');
    } finally {
      setLoading(false);
    }
  };

  const fetchTrainings = async () => {
    if (id) {
      try {
        const data = await getTrainingsByAgentId(Number(id));
        setTrainings(data);
        if (data.length > 0) {
          setCurrentTraining(data[0]);
        }
      } catch (error) {
        message.error('获取训练列表失败');
      }
    }
  };

  const fetchExamples = async () => {
    if (currentTraining?.id) {
      try {
        const data = await getExamplesByTrainingId(currentTraining.id);
        setExamples(data);
      } catch (error) {
        message.error('获取训练样例失败');
      }
    }
  };

  const fetchSessions = async () => {
    if (currentTraining?.id) {
      try {
        const data = await getSessionsByTrainingId(currentTraining.id);
        setTrainingSessions(data);
        
        // 如果有会话，加载最新会话的日志
        if (data.length > 0) {
          const latestSession = data[0];
          fetchSessionLogs(latestSession.id as number);
        }
      } catch (error) {
        message.error('获取训练会话失败');
      }
    }
  };

  const fetchSessionLogs = async (sessionId: number) => {
    try {
      const data = await getLogsBySessionId(sessionId);
      setTrainingLogs(data);
    } catch (error) {
      message.error('获取训练日志失败');
    }
  };

  const handleCreateTraining = async (values: TrainingFormValues) => {
    if (agent) {
      try {
        setIsLoading(true);
        const newTraining: Training = {
          ...values,
          agentId: Number(id),
          originalPrompt: agent.systemPrompt
        };
        
        const result = await createTraining(newTraining);
        setTrainings([result, ...trainings]);
        setCurrentTraining(result);
        setIsTrainingModalVisible(false);
        form.resetFields();
        message.success('创建训练成功');
      } catch (error) {
        message.error('创建训练失败');
      } finally {
        setIsLoading(false);
      }
    }
  };

  const handleAddExample = async (values: TrainingExample) => {
    if (currentTraining?.id) {
      try {
        setIsLoading(true);
        const example = {
          ...values,
          trainingId: currentTraining.id
        };
        
        const result = await createExample(example);
        setExamples([...examples, result]);
        setIsExampleModalVisible(false);
        exampleForm.resetFields();
        message.success('添加样例成功');
      } catch (error) {
        message.error('添加样例失败');
      } finally {
        setIsLoading(false);
      }
    }
  };

  const handleDeleteExample = async (id: number) => {
    try {
      await deleteExample(id);
      setExamples(examples.filter(ex => ex.id !== id));
      message.success('删除样例成功');
    } catch (error) {
      message.error('删除样例失败');
    }
  };

  const handleStartTraining = async () => {
    if (currentTraining?.id) {
      try {
        setIsLoading(true);
        await startTraining(currentTraining.id);
        // 重新获取训练信息
        fetchTrainings();
        message.success('开始训练成功');
      } catch (error) {
        message.error('开始训练失败');
      } finally {
        setIsLoading(false);
      }
    }
  };

  const handleApplyTraining = async () => {
    if (currentTraining?.id && agent?.id) {
      try {
        setIsLoading(true);
        await applyTrainingResult(currentTraining.id);
        
        // 更新智能体信息
        fetchAgentData();
        message.success('应用训练结果成功');
      } catch (error) {
        message.error('应用训练结果失败');
      } finally {
        setIsLoading(false);
      }
    }
  };

  const exampleColumns = [
    {
      title: '输入',
      dataIndex: 'input',
      key: 'input',
      width: '40%',
      ellipsis: true,
    },
    {
      title: '期望输出',
      dataIndex: 'expectedOutput',
      key: 'expectedOutput',
      width: '40%',
      ellipsis: true,
    },
    {
      title: '操作',
      key: 'action',
      render: (_: any, record: TrainingExample) => (
        <Button 
          type="text" 
          danger 
          icon={<DeleteOutlined />} 
          onClick={() => record.id && handleDeleteExample(record.id)}
        />
      ),
    },
  ];

  const getStatusBadge = (status?: string) => {
    if (!status) return <Badge status="default" text="未开始" />;
    
    const statusMap: Record<string, { status: "success" | "processing" | "default" | "error" | "warning", text: string }> = {
      pending: { status: 'default', text: '未开始' },
      running: { status: 'processing', text: '训练中' },
      completed: { status: 'success', text: '已完成' },
      failed: { status: 'error', text: '失败' },
      paused: { status: 'warning', text: '已暂停' },
    };
    
    const config = statusMap[status] || { status: 'default', text: status };
    return <Badge status={config.status} text={config.text} />;
  };

  return (
    <div className="agent-train-container">
      <div className="agent-train-header">
        <Space>
          <Button 
            icon={<ArrowLeftOutlined />} 
            onClick={() => navigate('/agent')}
          >
            返回
          </Button>
          <Title level={2}>智能体训练</Title>
        </Space>
        
        <Button
          type="primary"
          icon={<PlusOutlined />}
          onClick={() => setIsTrainingModalVisible(true)}
        >
          创建训练
        </Button>
      </div>

      {agent && (
        <Layout className="agent-train-layout">
          <Sider width={300} className="agent-train-sider">
            <Card title="智能体信息" bordered={false}>
              <Descriptions column={1} size="small">
                <Descriptions.Item label="名称">{agent.name}</Descriptions.Item>
                <Descriptions.Item label="版本">{agent.version || 1}</Descriptions.Item>
                <Descriptions.Item label="可训练">{agent.trainable ? '是' : '否'}</Descriptions.Item>
              </Descriptions>
              
              <Divider orientation="left">系统提示词</Divider>
              <div className="system-prompt-container">
                <Paragraph ellipsis={{ rows: 6, expandable: true }}>
                  {agent.systemPrompt}
                </Paragraph>
              </div>
              
              <Divider orientation="left">训练列表</Divider>
              {trainings.length > 0 ? (
                <div className="training-list">
                  {trainings.map(training => (
                    <div 
                      key={training.id} 
                      className={`training-item ${currentTraining?.id === training.id ? 'active' : ''}`}
                      onClick={() => setCurrentTraining(training)}
                    >
                      <div className="training-item-title">
                        <span>{training.name}</span>
                        {getStatusBadge(training.status)}
                      </div>
                      <div className="training-item-desc">
                        迭代: {training.currentIteration || 0}/{training.maxIterations}
                      </div>
                    </div>
                  ))}
                </div>
              ) : (
                <Text type="secondary">暂无训练数据</Text>
              )}
            </Card>
          </Sider>
          
          <Content className="agent-train-content">
            {currentTraining ? (
              <>
                <div className="training-header">
                  <Title level={4}>{currentTraining.name}</Title>
                  <div className="training-actions">
                    <Space>
                      <Button 
                        type="primary" 
                        icon={<PlayCircleOutlined />}
                        onClick={handleStartTraining}
                        loading={isLoading}
                        disabled={examples.length === 0 || currentTraining.status === 'running'}
                      >
                        开始训练
                      </Button>
                      <Button 
                        icon={<SaveOutlined />}
                        onClick={handleApplyTraining}
                        disabled={!currentTraining.optimizedPrompt || currentTraining.status !== 'completed'}
                      >
                        应用结果
                      </Button>
                    </Space>
                  </div>
                </div>
                
                <Card 
                  title={
                    <div className="card-title-with-action">
                      <span>训练样例</span>
                      <Button 
                        type="primary" 
                        size="small" 
                        icon={<PlusOutlined />}
                        onClick={() => setIsExampleModalVisible(true)}
                      >
                        添加样例
                      </Button>
                    </div>
                  } 
                  className="training-examples-card"
                >
                  <Table
                    dataSource={examples}
                    columns={exampleColumns}
                    rowKey="id"
                    pagination={false}
                    size="small"
                    scroll={{ y: 250 }}
                  />
                </Card>
                
                <Card title="训练日志" className="training-log-card">
                  {currentTraining.status === 'running' && (
                    <Progress 
                      percent={Math.round((currentTraining.currentIteration || 0) / currentTraining.maxIterations * 100)} 
                      status="active" 
                    />
                  )}
                  
                  {trainingSessions.length > 0 ? (
                    <div className="training-sessions">
                      <div className="session-tabs">
                        {trainingSessions.map(session => (
                          <Badge 
                            key={session.id} 
                            count={session.iteration} 
                            className={`session-tab ${trainingLogs[0]?.sessionId === session.id ? 'active' : ''}`}
                            onClick={() => session.id && fetchSessionLogs(session.id)}
                          >
                            轮次 {session.iteration}
                          </Badge>
                        ))}
                      </div>
                      
                      <div className="training-logs">
                        {trainingLogs.map(log => (
                          <div key={log.id} className="log-item">
                            <div className="log-content">
                              <div className="log-header">
                                <div>样例 #{log.exampleId}</div>
                                <div>相似度: {(log.similarity || 0) * 100}%</div>
                              </div>
                              <div className="log-body">
                                <div className="log-input">
                                  <Text type="secondary">输入: </Text>
                                  <Text>{log.input}</Text>
                                </div>
                                <div className="log-expected">
                                  <Text type="secondary">期望: </Text>
                                  <Text>{log.expectedOutput}</Text>
                                </div>
                                <div className="log-actual">
                                  <Text type="secondary">实际: </Text>
                                  <Text>{log.actualOutput}</Text>
                                </div>
                              </div>
                            </div>
                          </div>
                        ))}
                      </div>
                    </div>
                  ) : (
                    <div className="empty-log">
                      {examples.length > 0 ? (
                        <Text type="secondary">点击"开始训练"按钮开始训练过程</Text>
                      ) : (
                        <Text type="secondary">请先添加训练样例</Text>
                      )}
                    </div>
                  )}
                </Card>
              </>
            ) : (
              <div className="empty-training">
                <Text type="secondary">请选择或创建一个训练任务</Text>
              </div>
            )}
          </Content>
        </Layout>
      )}
      
      <Modal
        title="创建训练"
        open={isTrainingModalVisible}
        onCancel={() => setIsTrainingModalVisible(false)}
        footer={null}
      >
        <Form
          form={form}
          layout="vertical"
          onFinish={handleCreateTraining}
        >
          <Form.Item
            name="name"
            label="训练名称"
            rules={[{ required: true, message: '请输入训练名称' }]}
          >
            <Input placeholder="请输入训练名称" />
          </Form.Item>
          
          <Form.Item
            name="description"
            label="训练描述"
          >
            <TextArea placeholder="请输入训练描述" rows={3} />
          </Form.Item>
          
          <Form.Item
            name="maxIterations"
            label="最大迭代次数"
            initialValue={5}
            rules={[{ required: true, message: '请输入最大迭代次数' }]}
          >
            <Input type="number" min={1} max={20} />
          </Form.Item>
          
          <Form.Item>
            <div className="modal-footer">
              <Button onClick={() => setIsTrainingModalVisible(false)}>取消</Button>
              <Button type="primary" htmlType="submit" loading={isLoading}>创建</Button>
            </div>
          </Form.Item>
        </Form>
      </Modal>
      
      <Modal
        title="添加训练样例"
        open={isExampleModalVisible}
        onCancel={() => setIsExampleModalVisible(false)}
        footer={null}
      >
        <Form
          form={exampleForm}
          layout="vertical"
          onFinish={handleAddExample}
        >
          <Form.Item
            name="input"
            label="输入内容"
            rules={[{ required: true, message: '请输入样例输入内容' }]}
          >
            <TextArea placeholder="请输入样例输入内容" rows={4} />
          </Form.Item>
          
          <Form.Item
            name="expectedOutput"
            label="期望输出"
            rules={[{ required: true, message: '请输入期望输出内容' }]}
          >
            <TextArea placeholder="请输入期望输出内容" rows={4} />
          </Form.Item>
          
          <Form.Item>
            <div className="modal-footer">
              <Button onClick={() => setIsExampleModalVisible(false)}>取消</Button>
              <Button type="primary" htmlType="submit" loading={isLoading}>添加</Button>
            </div>
          </Form.Item>
        </Form>
      </Modal>
    </div>
  );
};

export default AgentTrain; 